# -*- coding: utf-8 -*-
# time: 2025/4/9 14:52
# file: llm_ts.py
# author: hanson
"""

测试微调后的模型
"""
from llama_index.llms.huggingface import HuggingFaceLLM

# 加载微调后的模型

llm = HuggingFaceLLM(
    model_name=r"./qwen2-0.5b-finetuned",
    tokenizer_name=r"./qwen2-0.5b-finetuned",
    model_kwargs={"trust_remote_code": True},
    tokenizer_kwargs={"trust_remote_code": True}
)

# 创建对话链
from langchain_core.prompts import PromptTemplate

prompt = PromptTemplate.from_template(
    """基于以下指令生成回答：
    Instruction: {instruction}
    Input: {input}
    Answer:"""
)

chain = prompt | llm

# 测试推理
result = chain.invoke({
    "instruction": "解释机器学习",
    "input": "用简单的中文说明"
})
print(result["text"])